import json
import os
from glob import glob
from pathlib import Path

from option import parse_args


def get_original_video_paths(root_dir, basename=False):
    originals = set()
    originals_v = set()
    for json_path in glob(os.path.join(root_dir, "dfdc_train_part_*/metadata.json")):
        dir = Path(json_path).parent
        with open(json_path, "r") as f:
            metadata = json.load(f)
        for k, v in metadata.items():
            original = v.get("original", None)
            if v["label"] == "REAL":
                original = k
                originals_v.add(original)
                originals.add(os.path.join(dir, original))
    originals = list(originals)
    originals_v = list(originals_v)
    return originals_v if basename else originals


def main():
    args = parse_args()
    originals = get_original_video_paths(args.root_dir)
    print(len(originals))  # 4743


if __name__ == '__main__':
    main()
